# Copyright 2022 The XFL Authors. All rights reserved.
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from typing import Union

from .drbg_base import DRBGBase
from .hmac_drbg import HMAC_DRBG

#######################################################
#    NIST SP 800-90A
#    NIST SP 800-57 PART 1 Table 2, Table 3, Table 4
#######################################################


def get_drbg_inst(name: str,
                  entropy: Union[bytes, bytearray],
                  method: str = '',
                  nonce: Union[bytes, bytearray] = b'',
                  additional_data: Union[bytes, bytearray] = b'') -> DRBGBase:
    """"
    Get DRBG instance to generate random numbers by a given entropy. 
    Reference material is NIST SP 800-90A[1].

    Args:
        name (str): support "hmac_drbg", ("hash_drbg" and "ctr_drbg" in future).
        entropy (Union[bytes, bytearray]): see [1]. For example, entropy can be a truncated secret key generated by 
            diffie-hellman key exchange.
        method (str, optional): depend on name. Defaults to ''. see Specificatioins.
        nonce (Union[bytes, bytearray], optional): see [1]. Defaults to b''.
        additional_data (Union[bytes, bytearray], optional): see [1]. Defaults to b''.

    Returns:
        DRBGBase: a DRBG instance
        
    Specifications:
        for name = 'hmac_drbg', method can be one of the ['sha1', 'sha224', 'sha256', 'sha384', 'sha512', ('sm3' in future)],
        if not specified, default method is 'sha512'.
        Length(number of bytes) of the entropy should satisfy:
        
            method    min_entropy_length  max_entropy_length
            sha1      10                  1 << 32
            sha224    14                  1 << 32
            sha256    16                  1 << 32
            sha384    24                  1 << 32
            sha512    32                  1 << 32
            sm3       16                  1 << 32
            
        Also, length of additional_data should less equal than 1 << 32.
    
    """
    
    opt = {
        "hmac_drbg": {
            'cls': HMAC_DRBG,
            'method': 'sha512' if len(method) == 0 else method
        }
    }
    
    try:
        v = opt[name]
        return v['cls'](v['method'], entropy, nonce, additional_data)
    except KeyError:
        raise KeyError(f"{name} not supported in drbg algorithm, use one of {list(opt.keys())} instead.")
    except Exception as e:
        raise e
